package com.sean.leetcode.test;
import java.util.*;

// 平衡二叉树节点类
class TreeNode{
    public int val;
    public int bf;          // 平衡因子: 左子树高 -1 相等 0 右子树高 1
    public TreeNode left;
    public TreeNode right;
    public TreeNode parent; // 父节点

    public TreeNode() {
    }

    public TreeNode(int val) {
        this.val = val;
    }
}

/********************
 * 平衡调整步骤
 * 1. 找平衡因子 = 2
 * 2. 找插入新节点后，失去平衡的最小子树
 *    子树的要求：
 *      a. 距离新插入的节点最近
 *      b. 平衡因子绝对值大于 1 的节点作为最小子树的根
 * 3. 平衡调整
 *      a. LL 型 ———— R
 *          中为支，高右转（将高的节点顺时针旋转）
 *      b. RR 型 ———— L
 *          中为支，高左转（将高的节点逆时针旋转）
 *      c. LR 型 ———— LR
 *          下二整体先左转（逆时针旋转），后与 LL 同
 *      d. RL 型 ———— RL
 *          下二整体先右转（顺时针旋转），后与 RR 同
 ********************/
public class AVLTree {
    private TreeNode root;

    public boolean insert(int val){
        // 插入
        TreeNode node = new TreeNode(val);
        if(this.root == null){          // 根节点为空
            root = node;                // 新插入节点直接作为根节点
            return true;
        }
        // 根据二叉搜索树性质，找到适合插入的位置
        TreeNode parent = null;
        TreeNode cur = this.root;
        while(cur != null){
            if(cur.val > val){
                parent = cur;
                cur = cur.left;
            }else if(cur.val < val){
                parent = cur;
                cur = cur.right;
            }else{                      // 该节点已经存在不允许插入
                return false;
            }
        }
        if(parent.val > val){
            parent.left = node;
        }else{
            parent.right = node;
        }
        node.parent = parent;
        /* 调整平衡因子
         * 从新插入节点的父节点开始调整，如果父节点的平衡因子变成了0，则停止调整
         * 如果该节点的父节点的平衡因子不是0，则继续向上调整，
         * 直到某个节点的平衡因子变成0或者调整到了根节点。
         * 调整平衡因子的策略是，如果插入的节点是该节点的左孩子节点则平衡因子-1
         * 如果是该节点的右孩子节点，则平衡因子+1
         */
        cur = node;
        while(parent != null){
            if(cur == parent.left){                 // 新插入节点是父节点的左子树，那么父节点的平衡因子减 1
                parent.bf --;
            }else{
                parent.bf ++;
            }

            if(parent.bf == 0){                     // 判断平衡因子，为 0 说明树平衡，则退出
                break;
            }else if(Math.abs(parent.bf) <= 1){     // 不为 0 则继续向上调整
                cur = parent;
                parent = parent.parent;
            }else{                                  // 平衡因子绝对值大于等于 2 了，已经不是平衡二叉树了，需要进行旋转
                if(parent.bf == -2){                // 左子树更高
                    if(parent.left.bf == -1){       // 左子树的左子树更高
                        // LL 型
                        this.RotateR(parent);
                        break;
                    }else{                          // 左子树的右子树高
                        // LR 型
                        this.RotateLR(parent);
                        break;
                    }
                }else if(parent.bf == 2){           // 右子树更高
                    if(parent.right.bf == 1){       // 右子树的右子树更高
                        // RR 型
                        this.RotateL(parent);
                        break;
                    }else{                          // 右子树的左子树更高
                        // RL 型
                        this.RotateRL(parent);
                        break;
                    }
                }
            }
        }
        return true;
    }
    // 右单旋
    private void RotateR(TreeNode parent){
        /*
         *          4（parent）
         *      2
         *  1       3(null)
         */
        TreeNode parent_left = parent.left;              // 不平衡节点的左子树   2
        TreeNode parent_left_right = parent_left.right;  // 不平衡节点的左子树的右子树 3
        TreeNode parent_parent = parent.parent;          // 不平衡节点的父节点

        // 4.left 指向 3，parent节点的左指针指向左子树的右节点
        parent.left = parent_left_right;
        // 4.parent 指向 2
        parent.parent = parent_left;
        // 如果 2 的右子树 3 存在
        if(parent_left_right != null){
            // 更新 3.parent = 4
            parent_left_right.parent = parent;
        }
        // 更新 2 的父节点和右节点
        parent_left.parent = parent_parent;         // 2.parent = 4 的 parent
        parent_left.right = parent;                 // 2.right = 4
        if(parent_parent != null){                  // 如果 4 的 parent 存在，说明 4 是某个节点的子树
            if(parent_parent.left == parent){       // 如果 4 是某个节点的左子树
                parent_parent.left = parent_left;   // 将该节点的左子树改为 2
            }else{
                parent_parent.right = parent_left;  // 改为 2
            }
        }else{  // 4 的 parent 不存在，说明 4 是 root
            this.root = parent_left; // root 改为 2
        }
        // 平衡因子置 0
        parent.bf = 0;
        parent_left.bf = 0;
    }

    // 左单旋
    private void RotateL(TreeNode parent){
        /*
         *          1
         *              3
         *      2(null)     4
         */
        TreeNode parent_parent = parent.parent;
        TreeNode parent_right = parent.right;
        TreeNode parent_right_left = parent_right.left;
        // 处理 parent
        parent.right = parent_right_left;
        parent.parent = parent_right;
        // 判断右子树的左子树是否存在
        if(parent_right_left != null){
            parent_right_left.parent = parent;
        }
        // 更新右子树
        parent_right.left = parent;
        parent_right.parent = parent_parent;
        // 判断原 parent 的 parent 是否存在
        if(parent_parent != null){
            // 是某个节点的子树
            if(parent_parent.left == parent){
                parent_parent.left = parent_right;
            }else{
                parent_parent.right = parent_right;
            }
        }else{
            this.root = parent_right;
        }
        parent.bf = 0;
        parent_right.bf = 0;
    }

    // 左右旋转
    private void RotateLR(TreeNode parent){
        /* 插入 8 导致 9 的因子变为 2
         *      9（parent）
         *    6     11
         *  0   7
         *        8
         *
         */
        TreeNode parent_left = parent.left;
        // bf 用于判断 6 的右子树的平衡因子（1）
        int bf = parent_left.right.bf;
        // 左旋左子树（6）
        this.RotateL(parent_left);
        /*  左旋之后变为
         *          9 (parent)
         *       7      11
         *    6     8
         *  0
         */
        // 右旋（9）
        this.RotateR(parent);
        /*  右旋之后变为
         *          7
         *       6      9 (parent)
         *     0       8   11
         */
        if(bf == 1) {
            // 说明 parent 接收了 6 的右子树的右子树（8）
            parent_left.bf = -1;
        }else if(bf == -1){
            // 说明 6 接收了 6 的右子树的左子树
            /*
             *       7
             *    6       9(parent)
             *  0  hear      11
             */
            parent.bf = 1;
        }
    }

    // 右左旋转
    private void RotateRL(TreeNode parent){
        TreeNode parent_right = parent.right;
        int bf = parent_right.left.bf;
        // 右旋
        this.RotateR(parent_right);
        // 左旋
        this.RotateL(parent);
        // 调整平衡因子
        if(-1 == bf){
            parent_right.bf = 1;
        }else if(1 == bf){
            parent.bf = -1;
        }
    }

    // 层序遍历打印二叉树
    public void printTree(){
        final List<List<Integer>> ans = this.traversal();
        for (List<Integer> tmp : ans) {
            for(int i = 0; i < tmp.size(); i++){
                if(i == 0){
                    System.out.print("[");
                }
                if(tmp.get(i) != null && i < tmp.size() - 1){
                    System.out.print(tmp.get(i) + ", ");
                }else if(tmp.get(i) != null && i == tmp.size() - 1){
                    System.out.print(tmp.get(i));
                }else if(tmp.get(i) == null && i < tmp.size() - 1){
                    System.out.print("x, ");
                }else{
                    System.out.print("x");
                }
                if(i == tmp.size() - 1){
                    System.out.println("]");
                }
            }
        }
    }
    // 层序遍历
    private List<List<Integer>> traversal(){
        List<List<Integer>> ans = new ArrayList<>();
        Queue<TreeNode> queue = new LinkedList<>();
        queue.offer(root);
        while(queue.size() > 0){
            List<Integer> tmp = new ArrayList<>();
            int size = queue.size();
            for(int i = 0; i < size; i++){
                TreeNode cur = queue.poll();
                if(cur != null){
                    tmp.add(cur.val);
                }else{
                    tmp.add(null);
                    continue;
                }
                queue.offer(cur.left);
                queue.offer(cur.right);
            }
            ans.add(tmp);
        }
        return ans;
    }
}

class Test{
    public static void main(String[] args) {
        AVLTree avlTree = new AVLTree();
        // test LL
        avlTree.insert(4);
        avlTree.insert(2);
        avlTree.insert(1);
        avlTree.insert(0);
        avlTree.insert(-1);
        avlTree.insert(-2);
        // test RR
        avlTree.insert(-2);
        avlTree.insert(-1);
        avlTree.insert(0);
        avlTree.insert(1);
        avlTree.insert(2);
        avlTree.insert(4);
        // test LR
        avlTree.insert(9);
        avlTree.insert(6);
        avlTree.insert(11);
        avlTree.insert(0);
        avlTree.insert(7);
        avlTree.insert(8);  // 发生 LR
        // test RL
        avlTree.insert(2);
        avlTree.insert(1);
        avlTree.insert(5);
        avlTree.insert(4);
        avlTree.insert(6);
        avlTree.insert(3);  // 发生 RL
        avlTree.printTree();
    }
}


